// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
//

#include <kllm/kai/proto_helper.h>
#include <kllm/utility/json.h>
#include <kllm/core/task_slot.h>

namespace kllm {

    bool is_string(const KaiPrompts &prompts) {
        if (prompts.values_size() == 1) {
            return prompts.values(0).has_string_value();
        }
        return false;
    }

    bool is_number(const KaiPrompts &prompts) {
        if (prompts.values_size() == 1) {
            return prompts.values(0).has_number_value();
        }
        return false;
    }

    bool is_string_number_mix(const KaiPrompts &prompts) {
        bool is_str = false;
        bool is_num = false;
        for (int i = 0; i < prompts.values_size(); ++i) {
            if (prompts.values(i).has_string_value()) {
                is_str = true;
            }
            if (prompts.values(i).has_number_value()) {
                is_num = true;
            }
            if (prompts.values(i).has_list_value()) {
                return false;
            }
        }
        return is_str && is_num;
    }

    bool is_list(const KaiPrompts &prompts) {
        return prompts.values_size() > 1;
    }

    bool is_string_list(const KaiPrompts &prompts) {
        if (prompts.values_size() <= 1) {
            return false;
        }
        for (int i = 0; i < prompts.values_size(); ++i) {
            if (!prompts.values(i).has_string_value()) {
                return false;
            }
        }
        return true;
    }

    bool is_number_list(const KaiPrompts &prompts) {
        if (prompts.values_size() <= 1) {
            return false;
        }
        for (int i = 0; i < prompts.values_size(); ++i) {
            if (!prompts.values(i).has_number_value()) {
                return false;
            }
        }
        return true;
    }

    bool has_list_value(const KaiPrompts &prompts) {
        for (int i = 0; i < prompts.values_size(); ++i) {
            if (prompts.values(i).has_list_value()) {
                return true;
            }
        }
        return false;
    }

    bool extract_number_list(const KaiPrompts &prompts, std::vector<int32_t> &ret) {
        ret.reserve(prompts.values_size());
        for (size_t i = 0; i < prompts.values_size(); i++) {
            if (prompts.values(i).has_number_value()) {
                return false;
            }
            ret.push_back(prompts.values(i).number_value());
        }
        return true;
    }

    bool extract_number_list(const PromptsValue &prompts, std::vector<int32_t> &ret) {
        if (prompts.has_string_value()) {
            return false;
        }
        if (prompts.has_number_value()) {
            ret = {prompts.number_value()};
            return true;
        }

        if (prompts.has_list_value()) {
            ret.reserve(prompts.list_value().values_size());
            for (auto &it: prompts.list_value().values()) {
                if (!it.has_number_value()) {
                    return false;
                }
                ret.push_back(it.number_value());
            }
        }
        return true;
    }

    PromptsValueType detect_prompt_type(const PromptsValue &prompts) {
        if (prompts.has_number_value()) {
            return TOKEN_TYPE;
        }
        if (prompts.has_string_value()) {
            return STRING_TYPE;
        }

        if (prompts.has_list_value()) {
            return LIST_TYPE;
        }

        return EMPTY_TYPE;
    }

    PromptsValueType detect_prompt_list_type(const PromptsValue &prompts) {
        bool is_str = false;
        bool is_num = false;
        bool is_list = false;
        for (int i = 0; i < prompts.list_value().values_size(); ++i) {
            if (prompts.list_value().values(i).has_string_value()) {
                is_str = true;
            }
            if (prompts.list_value().values(i).has_number_value()) {
                is_num = true;
            }
            if (prompts.list_value().values(i).has_list_value()) {
                is_list = true;
            }
        }
        if (is_list) {
            return LIST_TYPE;
        }
        if (is_str && is_num) {
            return MIX_LIST_TYPE;
        }
        if (is_str) {
            return STRING_LIST_TYPE;
        }
        if (is_num) {
            return TOKEN_TYPE;
        }
        return EMPTY_TYPE;
    }

    KaiPromptsType detect_prompt_type(const KaiPrompts &prompts) {
        KaiPromptsType type;
        type.set_valid(true);
        // single
        if (prompts.values_size() == 1) {
            if (prompts.values(0).has_number_value()) {
                type.set_root_type(TOKEN_TYPE);
            } else if (prompts.values(0).has_string_value()) {
                type.set_root_type(STRING_TYPE);
            } else if (prompts.values(0).has_list_value()) {
                type.set_root_type(LIST_TYPE);
                auto ty = detect_prompt_list_type(prompts.values(0));
                type.mutable_sub_type()->Add(ty);
            }
            return type;
        }
        type.set_root_type(LIST_TYPE);
        for (size_t i = 0; i < prompts.values_size(); i++) {
            auto ty = detect_prompt_list_type(prompts.values(i));
            type.mutable_sub_type()->Add(ty);
        }
        return type;
    }

    // to [["string"],["string", 12, 54], [49]]
    bool flatten_to_mix_lists(const KaiPrompts &prompts, std::vector<PromptsValue> &ret) {
        bool see_list = false;
        for (auto &it: prompts.values()) {
            if (it.has_list_value()) {
                see_list = true;
                break;
            }
        }

        if (!see_list) {
            for (auto &it: prompts.values()) {
                PromptsValue seg;
                *seg.mutable_list_value()->add_values() = it;
                ret.push_back(seg);
            }
        } else {
            for (auto &it: prompts.values()) {
                // check no sub list
                if (it.has_list_value()) {
                    for (auto &lit: it.list_value().values()) {
                        if (lit.has_list_value()) {
                            return false;
                        }
                    }
                }
                ret.push_back(it);
            }
        }
        return true;
    }
}  // namespace kllm
